#include "VSKptTracker.hpp"


// uncomment following line to output result to cmd line
//#define DISPLAY


VSKptTracker::VSKptTracker()
{

}


VSKptTracker::VSKptTracker(string frameName)
{
	
	ref_img = imread(frameName.c_str());
	
	if (!ref_img.data)
	{
		printf("VSKptTracker: No ref image is found!! \n\n");
		getchar();
		//return -1;

	} else
	{
		imshow("VSKptTracker: reference image", ref_img);
		waitKey(0);
	}
	
	ClassInit();
}



int VSKptTracker::ClassInit()
{
	int ret = 1;

	termcrit = TermCriteria(CV_TERMCRIT_ITER|CV_TERMCRIT_EPS,20,0.03);
	subPixWinSize = Size(10,10);
	winSize = Size(31,31);

	MAX_COUNT = 500;

	// logic hidden/encapsulated in the class, supervising main program to use
	needToInit = 1;
	//nightMode = false;

	// init SIFT
	SIFTdect = SIFT(400, 3, 0.05, 10, 1.6);
	SIFTdect(ref_img, noArray(), keypoints_ref, descriptor_ref, false);   // get descriptors on reference image
	
	flag = 0; 

	return ret; 
}



VSKptTracker::VSKptTracker(Mat inputFrame)
{

	inputFrame.copyTo(ref_img);
  
	ClassInit();
}



VSKptTracker::~VSKptTracker()
{
  
  
}


// with returned pnts_ref and MATCH_MODE
int VSKptTracker::MatchOneIter(Mat& prevGray_in, Mat& gray_in, Mat& image, vector<Point2f> *points, int size, 
							   vector<Point2f>& error, vector<Point2f>& pnts_ref_in, MATCH_MODE mode)
{
	int ret = 1;

	if ( mode == MATCH_WITHOUT_POST)
	{
		MatchOneIter(prevGray_in, gray_in, image, points, size, error, pnts_ref_in);

	} else
		if ( mode == MATCH_WITH_POST)
		{
			MatchOneIter(prevGray_in, gray_in, image, points, size, error, pnts_ref_in);

			if (points[1].size() != pnts_ref_in.size())
			{
				printf("Error: points[1].sz and pnts_ref_in.sz after return\n");
				getchar();
			}

			OneIterPost(prevGray_in, gray_in, points, size, error);

		} else
		{
			printf("VSKptTracker: Match mode not supported!!");
			ret = 0;
		}

	return ret;
}



// overloading to return an additional var pnts_ref_in
int VSKptTracker::MatchOneIter(const Mat prevGray_in, const Mat gray_in, Mat& image, vector<Point2f> *points, int size, 
							   vector<Point2f>& error, vector<Point2f>& pnts_ref_in)
{
	int ret = 1;

	MatchOneIter(prevGray_in, gray_in, image, points, 3, error);	

	pnts_ref_in = pnts_ref;

	if (points[1].size() != pnts_ref_in.size())
		{
			printf("Error: points[1].sz and pnts_ref.sz after cut\n");
			getchar();
		}

	return ret;
}




int VSKptTracker::MatchOneIter(Mat& prevGray_in, Mat& gray_in, Mat& image, vector<Point2f> *points, int size, 
							   vector<Point2f>& error, MATCH_MODE mode)
{
	int ret = 1;

	if ( mode == MATCH_WITHOUT_POST)
	{
		MatchOneIter(prevGray_in, gray_in, image, points, size, error);

	} else
		if ( mode == MATCH_WITH_POST)
		{
			MatchOneIter(prevGray_in, gray_in, image, points, size, error);

			OneIterPost(prevGray_in, gray_in, points, 3, error);

		} else
		{
			printf("VSKptTracker: Match mode not supported!!");
			ret = 0;
		}

	return ret;
}




int VSKptTracker::MatchOneIter(const Mat prevGray_in, const Mat gray_in, Mat& image, vector<Point2f> *points, int size, 
							   vector<Point2f>& error)
{
  int ret = 1;

  if (needToInit == 1)
  {
	  printf("VSKptTracker: Error - please run InitMatching first!!\n");
	  ret = 0; 
  } else
  {
		prevGray_in.copyTo(prevGray);
		gray_in.copyTo(gray);
	  
		/*vector<uchar> status;
		vector<float> err;*/
		status.clear();
		err.clear();
		error.clear();

		if(prevGray.empty())
			gray.copyTo(prevGray);

		// get output matched points
		if (flag == OPTFLOW_USE_INITIAL_FLOW)
		{
			points[1] = points[0];
		}

		// LoggerIn << "before: " << points[0].size() << " ";

		calcOpticalFlowPyrLK(prevGray, gray, points[0], points[1], status, err, winSize,
			3, termcrit, flag, 0.001);
		
		
		checkConsistency(points[0], points[2], pnts_ref);

		if (points[0].size() != pnts_ref.size())
		{
			printf("Error: checkConsistency\n");
			getchar();
		}
		
		//LoggerIn << "Size check:";
		//LoggerIn << "Size check: " << points[0].size() << " " << points[1].size() << " " << err.size() << " " << status.size() << endl;
		//LoggerIn << "after: " << points[0].size() << endl;
		
		//fs << "err" << err; 
		
		for (int i=0; i<points[0].size(); i++)
		  if (status[i])
		  {
			// draw all the matching lines between 2 sequential frames
			Point2i pr = Point2i(cvRound(points[0][i].x),
						  cvRound(points[0][i].y));
			Point2i pi = Point2i(cvRound(points[1][i].x),
						  cvRound(points[1][i].y));
					
			// the vector pointing from prev to current frame
			line( image, pr, pi, CV_RGB(255,0,0),2 );
			circle( image, pr, 3, Scalar(0,255,0), -1, 8);  
			
			
			// draw all the matching lines between original and current frame
			pr = Point2i(cvRound(pnts_ref[i].x),
					  cvRound(pnts_ref[i].y));
			pi = Point2i(cvRound(points[1][i].x),
					  cvRound(points[1][i].y));
					
			// the vector pointing from prev to current frame
			line( image, pr, pi, CV_RGB(255,0,0),2 );
			circle( image, pr, 3, Scalar(0,255,0), -1, 8);
		    
			//p.x = matched_pnts_ref.at<Vec2f>(i)[0] - points[1][i].x;
			//p.y = matched_pnts_ref.at<Vec2f>(i)[1] - points[1][i].y;

			//error.push_back(p);
		  }
		
		// print output for checking 
	// 			for (int i=0; i<points[0].size(); i++)
	// 			  if (status[i])
	// 			  {
	// 				// draw all the matching lines
	// 				Point2i pr = Point2i(cvRound(points[0][i].x),
	// 						     cvRound(points[0][i].y));
	// 				printf("(%d,%d) ",pr.x,pr.y);
	// 			  }  
		  
	// 			printf("\n");
		  
	// 			for (int i=0; i<points[0].size(); i++)
	// 			  if (status[i])
	// 			  {
	// 				// draw all the matching lines
	// 				Point2i pi = Point2i(cvRound(points[1][i].x),
	// 						     cvRound(points[1][i].y));
	// 				printf("(%d,%d) ",pi.x,pi.y);
	// 			  }    
	// 			  
	// 			printf("\n");  

		// error calculation
		// remind: after this, error's size will be less than or equal to points[0]'s size
		//	   and it is an unknown number
		for (int i=0; i<points[0].size(); i++)
		  if (status[i] != 0)
		  {
			p.x = pnts_ref[i].x - points[1][i].x;
			p.y = pnts_ref[i].y - points[1][i].y;

			error.push_back(p);
		  }
		
#ifdef DISPLAY
		// calculate error at each iteration
		cout << "####### Before processing #######" << endl;
		cout << points[0].size() << " ";
		cout << points[1].size() << " ";
		cout << error.size() << " ";       // error's size is supposed to be less than or equal to points[0]'s size
		cout << status.size() << " ";
		cout << pnts_ref.size() << endl;
		//cout << matched_pnts_ref.cols << endl << endl;  
#endif
		
		size_t i, k;
		pnts_ref_tmp.clear();
		// Also resize pnts_ref accordingly
		for( i = k = 0; i < points[1].size(); i++ )   // what is this c++?
		{
		  
			// clicked point will also be added
			if( addRemovePt )
			{
				if( norm(point - points[1][i]) <= 5 )
				{
					addRemovePt = false;
					continue;
				}
			}
			
			// point w/o motion estimation is removed 
			if( !status[i] )
				continue;

			points[1][k++] = points[1][i];
			pnts_ref_tmp.push_back(pnts_ref[i]);
			//pnts_ref[k++] = pnts_ref[i];
			//circle( image, points[1][i], 3, Scalar(0,255,0), -1, 8);
		}
		points[1].resize(k);
		pnts_ref = pnts_ref_tmp;
		//printf("error here!\n");
		//pnts_ref.resize(k);
		//printf("error there!\n");

		if (points[1].size() != pnts_ref.size())
		{
			printf("Error: points[1].sz and pnts_ref.sz after cut\n");
			getchar();
		}
  
  }
  
  return ret;
}



int VSKptTracker::OneIterPost(Mat& prevGray_in, Mat& gray_in, vector<Point2f> *points, int size, vector<Point2f>& error)
{
  int ret = 1;
  
      prevGray_in.copyTo(prevGray);
      gray_in.copyTo(gray);
  
      // SubPix corner processing
      int addRemovePt = 0;    // following block is not used | except when mouse click is enabled
      if( addRemovePt && (points[1].size() < MAX_COUNT) )
      {
	      vector<Point2f> tmp;
	      tmp.push_back(point);
	      cornerSubPix( gray, tmp, winSize, cvSize(-1,-1), termcrit);
	      points[1].push_back(tmp[0]);
	      addRemovePt = false;
      }

      // log error output
      //fs << "Status" << status; 
      //fs << "Error" << error; 

#ifdef DISPLAY
      // copy current points/image to previous points/image
      // backup points[2]
      cout << "####### After processing #######" << endl; 
      cout << points[0].size() << " ";
      cout << points[1].size() << " ";    // points[1]'s size is supposed to be changed
      cout << error.size() << " ";        // error's size is supposed to be equal to points[1]'s size
      cout << status.size() << " ";
      cout << pnts_ref.size() << endl;
#endif
      
      points[2].clear();
      points[2] = points[1];   // points[2] = points[0] DEBUGGED 16/11

#ifdef DISPLAY
      cout << "####### 5 crucial vars #######" << endl; 
      printfDebuggingVars(pnts_ref, points[0], points[1], points[2], error, status);
#endif
      
      std::swap(points[1], points[0]);
      swap(prevGray, gray);
      
      prevGray.copyTo(prevGray_in);
      gray.copyTo(gray_in);
  
  return ret;
}


// initmatching with MATCH_MODE
int VSKptTracker::InitMatching(const Mat frame, Mat& prevGray_in, Mat& gray_in, vector<Point2f> *points, int size, vector<Point2f>& error, vector<Point2f>& pnts_ref_in, MATCH_MODE mode)
{
	int ret = 1;

	if (needToInit == 0)
	{
		printf("VSKptTracker: Error - cannot init matching again! Please restart program to do that!\n");

	} else
	{
	
		if (mode == MATCH_WITHOUT_POST)
		{
			InitMatching(frame, points, size, error);
			needToInit = 0;

		} else 
			if (mode == MATCH_WITH_POST)
			{
				InitMatching(frame, points, size, error);

				pnts_ref_in = points[0];
				
				OneIterPost(prevGray_in, gray_in, points, size, error);

				needToInit = 0;
			}
			else 
			{
				printf("VSKptTracker: Match mode not supported!!");
				ret = 0;
			}
	}

	return ret;
}



int VSKptTracker::InitMatching(const Mat frame, vector<Point2f> *points, int size, vector<Point2f>& error)
{
	int ret = 1;

	if (needToInit == 0)
	{
		printf("VSKptTracker: Error - cannot init matching again! Please restart program to do that!\n");

	} else
	{
			needToInit = 0;
			
			//keypoints_img, descriptor_img, enh_good_matches
			
			
			// restart vars
			ds.clear();
			error.clear();

			// automatic initialization
			// input: gray image, parameters e.g. 
			// output: points[1] (vector<Point2f>)
			// notes: no descriptor, just normal points
			// goodFeaturesToTrack(gray, points[1], MAX_COUNT, 0.01, 10, Mat(), 3, 0, 0.04);

			// from extracted points, get their corresponding sub-pixel corners
			// cornerSubPix(gray, points[1], subPixWinSize, Size(-1,-1), termcrit);
			
			// Perform initial match btw target and current		
			// Using 3rd version
			vote = callSIFTMatch(frame, descriptor_ref, keypoints_ref, ref_img,     // <-- used for webcam image
						keypoints_img, descriptor_img, enh_good_matches);

			// matches is used at every iteration to match current keypoints and reference keypoints
			matches = enh_good_matches;
			
			// log the 1st match
			// fs << "Matches";
			printf("######### INIT STEP ########\n");
			cout << "Good match Size=" << enh_good_matches.size() << endl;
			for (int i=0; i<enh_good_matches.size(); i++)
			{
				//fs << "Matches" << enh_good_matches[i].queryIdx << "\n";
				//cout << "(" << enh_good_matches[i].queryIdx << ",";
				//cout << enh_good_matches[i].trainIdx << ") ";
			}
			
			//getchar();
			// fs << "Matches" << enh_good_matches; 
			

			// Using original version
			//vote = callSIFTMatch(ref_img, frame);

			// using 2nd version
			//vote = callSIFTMatch(frame, descriptor_ref, keypoints_ref, ref_img);			

			// Feed the matched points to optical flow i.e. copy to points[1]
			XtractMatchingPoints(keypoints_ref, keypoints_img, enh_good_matches, matched_pnts_ref, matched_pnts_img);


			// get initial error
			pnts_ref.clear();
			for (int i=0; i<matched_pnts_ref.cols; i++)
			{
				// draw all the matching lines
				pr = Point2i(cvRound(matched_pnts_ref.at<Vec2f>(i)[0]),
						  cvRound(matched_pnts_ref.at<Vec2f>(i)[1]));
				pi = Point2i(cvRound(matched_pnts_img.at<Vec2f>(i)[0]),
						  cvRound(matched_pnts_img.at<Vec2f>(i)[1]));
							
				line( image, pr, pi, CV_RGB(255,0,0),2 );
				circle( image, pr, 3, Scalar(0,255,0), -1, 8);
				
				// get original reference points
				prf = Point2f(matched_pnts_ref.at<Vec2f>(i)[0],
						  matched_pnts_ref.at<Vec2f>(i)[1]);
				pnts_ref.push_back(prf);
				
				// error calculation
				pif = Point2f(matched_pnts_img.at<Vec2f>(i)[0],
						  matched_pnts_img.at<Vec2f>(i)[1]);
				p.x = prf.x - pif.x;
				p.y = prf.y - pif.y;
				error.push_back(p);
			}
			
			// get sub-pixel corners
			points[1].clear();
			for (int i=0; i<matched_pnts_img.cols; i++)   // get points[1] is duplicated with preceding code i.e. previous for statement
			{
				//points[1][i] = matched_pnts_img.at<Vec2f>(i);
				points[1].push_back(matched_pnts_img.at<Vec2f>(i));
			}
			
			// not using cornerSubPix at all
			// cornerSubPix(gray, points[1], subPixWinSize, Size(-1,-1), termcrit);
			
			// backup points[2]
			points[2].clear();
			points[0].clear();
			points[0] = pnts_ref;
			//points[2] = points[1];  // RISK?? -> YES! it's a RISK

	}
			
	return ret;
}




/**
 * @brief propagate pnts_ref from frame k-1 to frame k where No. keypoints in frame k-1 and frame k might be different 
 * 
 * @param points_a ...
 * @param points_b ...
 * @param pnts_ref ...
 * @return int
 */
int VSKptTracker::checkConsistency(vector<Point2f> points_a, vector<Point2f> points_b, vector<Point2f>& pnts_ref)
{
	int ret = 1;
	
	// check points[2] and points[0] consistency
	float tol = .05;
	vector<Point2f> vref_tmp;  // temp vector of ref 
	
	vector<int> flag(points_b.size(),0);
	
	for (int j=0; j<points_a.size(); j++)
	  for (int k=0; k<points_b.size(); k++)
	    if ((abs(points_a[j].x - points_b[k].x) <= tol)
	      &&(abs(points_a[j].y - points_b[k].y) <= tol)
	      &&(flag[k] == 0))
	    {
	      // coding
	      vref_tmp.push_back(pnts_ref[k]);	
	      flag[k] = 1;
	      break;
	    }
	
	pnts_ref = vref_tmp;
	
	return ret;
}



int VSKptTracker::printfDebuggingVars(vector<Point2f> pnts_ref, vector<Point2f> pnts_0, vector<Point2f> pnts_1, 
			vector<Point2f> pnts_2, vector<Point2f> error, vector<uchar> status)
{
	int ret = 1;

    Point2f pr;
    
    printf("reference points\n");
    for (int i=0; i<pnts_ref.size(); i++)
      {
	    // draw all the matching lines
	    Point2f pr = Point2f(pnts_ref[i].x,
				 pnts_ref[i].y);
	    printf("(%.1f,%.1f) ",pr.x,pr.y);
      }  
      
    printf("\n");
    
    
    printf("found points of prev keypoints in current frame\n");
    for (int i=0; i<pnts_0.size(); i++)
      {
	    // draw all the matching lines
	    Point2f pr = Point2f(pnts_0[i].x,
				 pnts_0[i].y);
	    printf("(%.1f,%.1f) ",pr.x,pr.y);
      }  
      
    printf("\n");
    
    
    printf("found matching points from motion in current frame\n");
    for (int i=0; i<pnts_1.size(); i++)
      {
	    // draw all the matching lines
	    Point2f pr = Point2f(pnts_1[i].x,
				 pnts_1[i].y);
	    printf("(%.1f,%.1f) ",pr.x,pr.y);
      }  
      
    printf("\n");
    
    
    printf("points[2]\n");
    for (int i=0; i<pnts_2.size(); i++)
      {
	    // draw all the matching lines
	    Point2f pr = Point2f(pnts_2[i].x,
				 pnts_2[i].y);
	    printf("(%.1f,%.1f) ",pr.x,pr.y);
      }  
      
    printf("\n");
  
    
    printf("error calculation for current frame\n");
    for (int i=0; i<error.size(); i++)
      {
	    // draw all the matching lines
	    Point2f pr = Point2f(error[i].x,
				 error[i].y);
	    printf("(%.1f,%.1f) ",pr.x,pr.y);
      }  
      
    printf("\n");
    
    
    printf("status for current frame\n");
    for (int i=0; i<status.size(); i++)
      {
	    printf("%d ",status[i]);
      }  
      
    printf("\n");

	return ret;
}